#include <mpi.h>
#include <stdio.h>
#include <limits.h>
#include <string.h>
#include <assert.h>
#include <memory.h>
#include <stdlib.h>


#define MPI_MERGEMESH_TAG 1


int main(void ) {
  int msg_size = INT_MAX-1;

    // Initialize the MPI environment
    MPI_Init(NULL, NULL);

    // Get the number of processes
    int nprocs;
    MPI_Comm_size(MPI_COMM_WORLD, &nprocs);

    // Get the rank of the process
    int myrank;
    MPI_Comm_rank(MPI_COMM_WORLD, &myrank);

    // Get the name of the processor
    char processor_name[MPI_MAX_PROCESSOR_NAME];
    int name_len;
    MPI_Get_processor_name(processor_name, &name_len);

    // From here, the code is copied from ParMmg and modified only to simulate
    // the gathering of a very large char array with gatherv function (if the
    // program is build with -DENABLE_GATHERV option) or MPI_Send/Recv ones
    size_t     pack_size_tot,next_disp;

    int        *rcv_pack_size;
    int        root,pack_size;
    char       *rcv_buffer,val,*buffer;

    INT_TYPE *displs,buf_idx;

    // Fake initialization of parameters and arrays on each proc
    root = 0;

    int8_t A = 65;
    if ( myrank == root ) {
      pack_size = 10;
      val = A;
    }
    else {
      int8_t rank_unit = myrank % 10;
      pack_size = msg_size;
      val = A + rank_unit;
    }

    /* Set buffer to suitable value */
    buffer = (char*)malloc(pack_size);
    rcv_pack_size = (int*)malloc(nprocs*sizeof(int));
    displs = (INT_TYPE*)malloc(nprocs*sizeof(INT_TYPE));

    memset(buffer,val,pack_size);

    printf(".. On rank %d, buffer is filled with %.5s\n",myrank,buffer);

    MPI_Gather(&pack_size,1,MPI_INT,rcv_pack_size,1,MPI_INT,root,MPI_COMM_WORLD);

    if ( myrank == root ) {
      displs[0] = 0;
      for ( int k=1; k<nprocs; ++k ) {

        next_disp = displs[k-1] + rcv_pack_size[k-1];
        if(next_disp>INT_MAX){
          /* The displacements argument to MPI_Gatherv() is an array of int
         * (signed) so the number of elements must be smaller than 2^31.
         * To get around this we must pack more data in a single element
         * or use multiple messages.
         */
          fprintf(stderr, "\n\n  ## Warning: INT MAX overflow in displs array \n\n");
          // MPI_Abort(MPI_COMM_WORLD, 1);      /* error detected only on root */
        }
        displs[k] = next_disp;
      }
      /* On root, we will gather all the meshes in rcv_buffer so we have to
       * compute the total pack size */
      pack_size_tot        = (size_t)(displs[nprocs-1])+(size_t)(rcv_pack_size[nprocs-1]);
      assert ( pack_size_tot < SIZE_MAX && "SIZE_MAX overflow" );

      /* root will write directly in the suitable position of rcv_buffer */
      buf_idx = displs[root];
    }
    else {
      /* on ranks other than root we just need to store the local mesh so buffer
       * will be of size pack_size */
      pack_size_tot = pack_size;
      /* we will write the mesh at the starting position */
      buf_idx = 0;
    }

    rcv_buffer = (char *)malloc(pack_size_tot);

 #ifndef ENABLE_GATHERV
    if ( myrank == root ) {
      /* With send/recv we have to copy buffer in rcv_buffer on rank root */
      memcpy(rcv_buffer,buffer,pack_size);

      for ( int i = 0; i < nprocs; ++i ) {
        if ( i != myrank ) {
          printf(" -- rank %d rcv %d data from %d at address %zu\n",root,rcv_pack_size[i],i,displs[i]);
          MPI_Recv(rcv_buffer + displs[i], rcv_pack_size[i], MPI_CHAR, i,
                   MPI_MERGEMESH_TAG, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
        }
      }
    } else {
      printf(" -- rank %d send a char array of size %d with value %c to %d\n",myrank,pack_size,buffer[0],root);
      MPI_Send(buffer, pack_size, MPI_CHAR, root, MPI_MERGEMESH_TAG,MPI_COMM_WORLD);
    }
#else
    MPI_Gatherv ( buffer,pack_size,MPI_CHAR,
                  rcv_buffer,rcv_pack_size,
                  displs,MPI_CHAR,root,MPI_COMM_WORLD );
#endif

    if ( myrank == root ) {
      printf("\n\n ==> Program succeed !!! gathered buffer %.5s ... %.5s ...\n",
             rcv_buffer,&rcv_buffer[pack_size_tot - rcv_pack_size[nprocs-1] - 2]);
    }

    free ( buffer );
    free ( rcv_buffer );

    // Finalize the MPI environment.
    MPI_Finalize();

    return 0;
}